"""pointnet test"""
import mindspore as ms
import numpy as np
from mind3d.models.blocks.t_net import STN3D, STNkD
from mindspore import Tensor


def test_stn3d(in_channel):
    """
    test the function of STN3D
    """
    data = Tensor(np.ones((32, 3, 1024)), ms.float32)
    trans = STN3D(in_channel=in_channel)
    out = trans(data)
    assert out.shape == (32, 3, 3), 'output shape not match'
    print("################################")
    print("STN3D test passed!")
    print("################################")


def test_stnkd(k):
    """
    test the function of STN3D
    """
    data = Tensor(np.ones((32, 64, 1024)), ms.float32)
    trans = STNkD(k=k)
    out = trans(data)
    assert out.shape == (32, 64, 64), 'output shape not match'
    print("################################")
    print("STNkD test passed!")
    print("################################")


if __name__ == "__main__":
    test_stn3d(in_channel=3)
    test_stnkd(k=64)
